from util.packagepretreat import read_pcap, get_raw,bytes2str_dict
from peer.peer import *
from dht.krcp import krcp
from uTP.uTP import utp_partion
from tracker.http_tracker import http_tracker
from tracker.udp_tracker import udp_tracker
from scapy.layers import http as h  # 该库将http也作为单独的一个数据包层：http层
from scapy.all import *
from util.logger import get_logger

pcap_fp = "D:\\c\\magnet-dht-master\\pcap\\tcp_dht_bt_test.pcap"# TODO 数据包改为实时捕获

analizied_list = []  # 最后结果存储[dictlist,head]
segment_dict = dict()  # 报文块存储字典{key = (ip.src,ip.dst), value = segment}
segment = bytes()  # 报文块
head = [0]  # 数据包的起始头位置
flag = dict()  # 字典记录当前特定方向的头{key = (ip.src,ip.dst), value = head:int}





i = 0



def tcp(data: bytes, i: int, segment_tuple: tuple):
    isbt = False
    #print(data)
    if data:
        ##print('bt')
        if segment_tuple in segment_dict:  # 存在前缀截断
            data = segment_dict[segment_tuple]+data
            # if i == 1758:
            #     ##print(len(data),data.hex())
            # #del segment_dict[segment_tuple]
            dictlist, segment, isbt = bt_partion(data)
            
            head[i] = flag[segment_tuple]  # 当前端的头记录
            if dictlist:  # 如果多个截断拼接完成，删除该方向的头
                del flag[segment_tuple]
                del segment_dict[segment_tuple]
        else:
            dictlist, segment, isbt = bt_partion(data)
        if segment:
            segment_dict[segment_tuple] = segment
            # 定义头的位置
            if segment_tuple not in flag:
                flag[segment_tuple] = i  # 记录特定方向的头
                head[i] = i  # 记录头的位置
        # ##print(head[i])
        # ##print(dictlist)
        return dictlist,isbt
    return None,isbt


"""
from scapy.layers import http as h#该库将http也作为单独的一个数据包层：http层

"""


def http(p) -> dict:
    try:
        field = None
        if p.haslayer(h.HTTPRequest):
            ##print('request')
            field = p[h.HTTPRequest].fields
            data = bytes.decode(field['Path'])
            result = http_tracker(data, True)

        elif p.haslayer(h.HTTPResponse):
            field = p[h.HTTPResponse].fields
            data = get_raw(p)
            result = http_tracker(data, False)
        field['URI Query'] = result
        return bytes2str_dict(field)
    except Exception as e:
        ##print(e)
        return None


def udp(data: bytes, i: int, segment_tuple):

    if data:
        dictlist = []
        prodict = krcp(data)  # DHT分析
        if prodict:
            kind = 'bt_dht'
            num_dht +=1
            return None,kind
        prodict, utp_len = utp_partion(data)  # utp分析
        if prodict:
            ##print('utp',prodict)
            kind = 'utp'
            num_utp +=1
            dictlist.append([prodict])
            if prodict['Type'] == '0':
                kind = 'bittorrent'
                num_bt +=1
                if segment_tuple in segment_dict:  # 存在前缀截断
                    data = segment_dict[segment_tuple]+data[utp_len:]
                    #del segment_dict[segment_tuple]
                    dictlist_temp, segment,_ = bt_partion(data)
                    head[i] = flag[segment_tuple]  # 当前端的头记录
                    if dictlist_temp:  # 如果多个截断拼接完成，删除该方向的头
                        del flag[segment_tuple]
                        del segment_dict[segment_tuple]
                else:
                    dictlist_temp, segment,_ = bt_partion(data[utp_len:])
                if segment:
                    segment_dict[segment_tuple] = segment
                    # 定义头的位置
                    if segment_tuple not in flag:
                        flag[segment_tuple] = i  # 记录特定方向的头
                        head[i] = i  # 记录头的位置
                if dictlist_temp:
                    dictlist = dictlist + dictlist_temp
            ##print(dictlist)
            return dictlist,kind
        else:
            return None,None
    return None,None


def analizy(package, i: int) -> list:

    '''
    输入：捕获的数据package，package的编号
    输出：字典列表，每一个字典存储解析的协议特征
    例如：
    [{'Type': '0', 'Version': '1', 'Extension': 0, 'Connection Id': 8000, 'Timestamp Microseconds': 425320392, 'Timestamp Difference Microseconds': 2295421527, 'Window Size': 1048576, 'Sequence Number': 42831, 'ACK Number': 56978}]
    '''
    kind = None
    global num_tracker
    head.append(-1)  # 当前包的头为自己，标记为-1
    result = http(package)
    ##print('http',result)
    if result:#tracker
        kind = 'http_tracker'
        return [result],-1,kind
    ip_layer = package.getlayer('IP')
    if not ip_layer:
        ip_layer = package.getlayer('IPv6')
    if not ip_layer:
        return None,-1,kind
    segment_tuple = (ip_layer.src, ip_layer.dst)
    data = get_raw(package)  # bytes数据
    if package.haslayer("UDP"):
        result = udp_tracker(data)
        ##print('udp_tracker', result)
        if result:
            kind = "udp_tracker"
            num_tracker += 1
            return result,-1,kind
        result,kind = udp(data, i, segment_tuple)
        ##print('udp',result)
    elif package.haslayer("TCP"):
        result,isbt = tcp(data, i, segment_tuple)
        if isbt or result:
            kind = 'bittorrent'
    return result,head[i],kind
 



def callback(package):
    
    global i
    ##print(i)
    dictlist = analizy(package,i)
    ##print('head',head[i])
    result = list(({'num': i}, dictlist, {'head': head[i]}))
    i+=1

def test():
    global i
    resultlsit = []
    pcap_fp = "D:\\c\\magnet-dht-master\\pcap\\udp_dht_bt_test.pcap"# TODO 数据包改为实时捕获
    packlist = read_pcap(pcap_fp)
    for p in packlist:
        i += 1
        print(i)
        dictlist,_,kind = analizy(p, i)
        print('dictlist',dictlist,kind)
        result = list(({'num': i}, dictlist, {'head': head[i]}))
        ##print(result)
        resultlsit.append(result)
        # analizied_list.append(DataAnalizied(dictlist,head[i]))
        if i > 100:
            break
    ##print(bt)
    ##print('num_bt:%d\nnum_dht:%d\nnum_tracker:%d\nnum_utp:%d\n' % (num_bt,num_dht,num_tracker,num_utp))

    


if __name__ == "__main__":
    test()
    #sniff(prn = callback,count = 0, iface = 'WLAN')

